import os
import os.path
import numpy as np
from matplotlib import pyplot as plt
from matplotlib import image as img
from heap import PriorityQueue

class State():
    def __init__(self, name=None, pos=None, cell=None, h=None, k=None, f="NEW", bp=None):
        self.name = name 
        self.pos = pos  #position
        self.cell = cell # cell type  Obstacle: 1, Empty: 0
        self.h = h 	    #h value
        self.k = k 	    #k value
        self.f = f	    #flag NEW、OPEN、CLOSED; also named tag
        self.bp = bp 	#backpointer
        # for d in dir(self):
        #     print(d)
        # raise Exception()
    
    def __str__(self):
        tmp_dict = self.__dict__
        tmp_dict["_bp_"] = id(self.bp)
        if self.bp:
            tmp_dict["_bp_pos"] = self.bp.pos
        else:
            tmp_dict["_bp_pos"] = None
        tmp_dict["_id"] = id(self)
        return str(tmp_dict)

    @property 
    def name(self):
        return self._name 

    @name.setter
    def name(self, name):
        self._name = name

    @property 
    def pos(self):
        return self._pos 

    @pos.setter
    def pos(self, pos):
        self._pos = pos

    @property 
    def cell(self):
        return self._cell

    @cell.setter 
    def cell(self, cell):
        self._cell = cell 
        
    @property 
    def h(self):
        return self._h 

    @h.setter
    def h(self, h):
        self._h = h 

    @property 
    def k(self):
        return self._k 

    @k.setter 
    def k(self, k):
        self._k = k

    @property 
    def f(self):
        return self._f 

    @f.setter
    def f(self, f):
        self._f = f 

    @property 
    def bp(self):
        return self._bp

    @bp.setter 
    def bp(self, bp):
        self._bp = bp


class DStarFinder():
    def __init__(self, start=None, goal=None, map_file=None, dynamic=None):
        self.start = start
        self.goal = goal
        self.load_base_map(map_file)
        self.dynamic = dynamic
        self.dynamic_flag = False

    @property
    def start(self):
        return self._start

    @start.setter
    def start(self, start):
        self._start = start

    @property
    def goal(self):
        return self._goal

    @goal.setter
    def goal(self, goal):
        self._goal = goal 
        # print(self._goal)

    @property
    def basemap(self):
        return self._basemap 

    @basemap.setter 
    def basemap(self, basemap):
        self._basemap = basemap

    @property
    def shape(self):
        return self._shape 

    @shape.setter 
    def shape(self, shape):
        self._shape = shape

    def load_base_map(self, map_file):
        # raise Exception()
        print(map_file)
        if os.path.exists(map_file):
            _temp_map = img.imread(map_file)
            _temp_map = _temp_map[:, :, 0]
            self.basemap = _temp_map
            self.shape = _temp_map.shape
            # print(self.shape)
            self.empty_val = 1  #cell type  Obstacle: 0, Empty: 1
            self.cost_obstacle = 10000
            self.cost_empty = 1
            self.cost_dynamic = 10000
            return True
        else:
            return False

    @property
    def graph(self):
        return self._graph

    @graph.setter
    def graph(self, graph):
        self._graph = graph

    def init_graph(self):
        if not self.basemap.shape:
            return False
        self.graph = np.full(self.shape, None)
        for row in range(self.shape[0]):
            for col in range(self.shape[1]):
                # print("row:{}, col:{}".format(row, col))
                _state = State()
                _state.name = "{}_{}".format(row, col)
                _state.pos = (row, col)
                _state.cell = self.basemap[row, col]  #cell type  Obstacle: 0, Empty: 1
                _state.h = float("inf")
                _state.k = None
                _state.f = 'NEW'
                _state.bp = None #backpointer
                self.graph[row, col] = _state
        return True

    @property
    def dynamic(self):
        return self._dynamic 

    @dynamic.setter
    def dynamic(self, dynamic):
        self._dynamic = dynamic

    @property
    def dynamic_flag(self):
        return self._dynamic_flag

    @dynamic_flag.setter
    def dynamic_flag(self, flag):
        self._dynamic_flag = flag

    def enable_dynamic(self):
        self.dynamic_flag = True

    def disable_dynamic(self):
        self.dynamic_flag = False
        
    @property
    def open(self):
        return self._open 
    
    @open.setter
    def open(self, open):
        self._open = open

    def open_enqueue(self, state):
        state.f = "OPEN"
        # print(state)
        self.open.enqueue(state)

    def open_dequeue(self):
        state = self.open.dequeue()
        state.f = "CLOSED"
        return state

    def open_update(self, state):
        if self.open.delete(state.pos):
            self.open.insert(state)

    def open_min(self):
        return self.open.head

    def _valid_pos(self, pos):
        if 0 <= pos[0] < self.shape[0] and 0 <= pos[1] < self.shape[1]:
            # print("valid pos: {}".format(pos))
            return True
        else:
            # print("invalid pos: {}".format(pos))
            return False

    def neighbors(self, state):
        nbs = []
        for i in [-1, 0, 1]:
            for j in [-1, 0, 1]:
                if i == 0 and j == 0:
                    continue
                new_pos = (state.pos[0]+i, state.pos[1]+j)
                # print(new_pos)
                if self._valid_pos(new_pos):
                    new_state = self.graph[new_pos]
                    # print(new_state.pos)
                    nbs.append(new_state)
        return nbs 

    @property
    def empty_val(self):
        return self._empty_val

    @empty_val.setter
    def empty_val(self, empty_val):
        self._empty_val = empty_val

    @property 
    def cost_obstacle(self):
        return self._cost_obstacle

    @cost_obstacle.setter
    def cost_obstacle(self, cost=10000):
        self._cost_obstacle = cost

    @property
    def cost_empty(self):
        return self._cost_empty

    @cost_empty.setter
    def cost_empty(self, cost=1):
        self._cost_empty = cost

    @property
    def cost_dynamic(self):
        return self._cost_dynamic

    @cost_dynamic.setter
    def cost_dynamic(self, cost=10000):
        self._cost_dynamic = cost

    def is_empty(self, state):
        if self._valid_pos(state.pos):
            # print("state.cell: {}  self.empty_val: {}".format(state.cell, self.empty_val))
            # raise Exception()
            if state.cell == self.empty_val:
                # raise Exception()
                return True
        return False

    def is_obstacle(self, state):
        return not self.is_empty(state)

    def has_dynamic(self, state):
        if self.dynamic_flag:
            if isinstance(self.dynamic, list):
                return state.pos in self.dynamic
        return False

    def arc_cost(self, from_u, to_v): #和原文有些不同为了从左向右阅读方便
        pos_dist = ((from_u.pos[0] - to_v.pos[0]) ** 2 + \
                    (from_u.pos[1] - to_v.pos[1]) ** 2) ** 0.5
        _cost = pos_dist * self.cost_empty
        if self.is_obstacle(from_u) or self.is_obstacle(to_v):
            _cost = self.cost_obstacle
        if self.has_dynamic(from_u) or self.has_dynamic(to_v):
            _cost = self.cost_dynamic
            # print("find dynamic from: {}".format(from_u))
            # print("find dynamic to: {}".format(from_u))
            # print("cost: {}".format(_cost))
        return _cost

    def _insert_open(self, state, h_val):
        if state.f == "NEW":
            state.k = h_val
        elif state.f == "OPEN":
            state.k = min(state.k, h_val)
        else: #"CLOSED"
            state.k = min(state.h, h_val)
        state.h = h_val
        if state.f == "NEW" or  state.f == "CLOSED":
            self.open_enqueue(state)
        else: #"CLOSED"
            self.open_update(state)
        return self.open_min()
    
    def _process_state(self, state=None):
        state = self.open_dequeue()
        if state is None:
            print("state is None")
            return None
        # print("state: {}".format(state))
        if state.k < state.h:#当遇到阻塞时才会触发 尝试提前修正
            print("in _process_state; state.k < state.h")
            print("state is: {}".format(state))
            neighbors = self.neighbors(state)
            for nb in neighbors:
                print("nb info: {}".format(nb))
                _arc_cost = self.arc_cost(state, nb)
                print("_arc_cost: {}".format(_arc_cost))
                if nb.h < state.k and state.h > _arc_cost + nb.h:
                    print("find nb.h < state.k: {}".format(state) )
                    state.bp = nb
                    state.h = _arc_cost + nb.h
                    print("after nb.h < state.k: {}".format(state) )
                    # raise Exception("find nb.h < state.k")
            # if state.pos == (39, 118):
            #     raise Exception()
        if state.k == state.h: #寻找最优路径
            # print("in _process_state; state.k == state.h")
            neighbors = self.neighbors(state)
            for nb in neighbors:
                print("{}'s neighbor {}".format(state.pos, nb))
                _arc_cost = self.arc_cost(nb, state)
                if nb.f == "NEW" or \
                    (nb.bp == state and nb.h != state.h + _arc_cost) or \
                    (nb.bp != state and nb.h > state.h + _arc_cost): 
                    nb.bp = state
                    print("insert: {}".format(nb.pos))
                    self._insert_open(nb, state.h + _arc_cost)
                    print("after insert: {}".format(nb))
        else:#修正新路径 
            print("state: {}".format(state))
            nbs = self.neighbors(state)
            for nb in nbs:
                print("st nb: {}".format(nb))
            # if state.pos == (39, 118):
                # raise Exception()
            # if state.k < state.h:
            neighbors = self.neighbors(state)
            for nb in neighbors:
                _arc_cost = self.arc_cost(nb, state)
                if nb.f == "NEW" or \
                   (nb.bp == state and nb.h != _arc_cost + state.h):
                   nb.bp = state
                   self._insert_open(nb, _arc_cost + state.h)
                else:
                    if nb.bp != state and nb.h > _arc_cost + state.h:
                        self._insert_open(state, state.h) #从raise模式，进入lower模式
                    else:
                        _arc_cost_2 = self.arc_cost(state, nb)
                        if nb.bp != state and state.h > _arc_cost_2 + nb.h and \
                            nb.f == "CLOSED" and nb.h > state.k: #次优解
                            self._insert_open(nb, nb.h) #通过 nb的扩展操作，将state重新优化
        return self.open_min()

    def block_dist(self, from_u, to_v):
        dist = abs(from_u.pos[0] - to_v.pos[0]) + abs(from_u.pos[1] - to_v.pos[1])
        return dist

    def _batch_modify_detected_cost(self, states):
        for state in states:
            nbs = self.neighbors(state)
            for nb in nbs:
                self._modify_cost(nb)
        
    def _batch_modify_visible_cost(self, from_u):
        for dy in self.dynamic:
            dy_state = self.graph[dy]
            if self.block_dist(from_u, dy_state) >= 20: #20街区距离都可以看见
                print(from_u)
                print(dy_state)
                raise Exception()
                continue 
            self._batch_modify_cost(dy_state)
        
    def _batch_modify_cost(self, to_v, from_u=None, val=None):
        self._modify_cost(to_v, from_u, val)
        nbs = self.neighbors(to_v)
        for nb in nbs:
            self._modify_cost(nb, from_u, val)

    def _modify_cost(self, to_v, from_u=None, val=None):
        #因为cost在底层是动态求解的，当存在动态障碍物时会自动修正cost的计算逻辑
        #所以modify_cost 实际未做任何操作，此处只是象征性表示。
        #使节点state进入RAISE模式
        # arc_cost_from_u_to_v = val #此处不做任何实现，因为cost在底层是动态求解的，当存在动态障碍物时会自动修正cost的计算逻辑
        if to_v.f == "CLOSED":
            self.open_enqueue(to_v) #把to_v重新入队，处于lower模式，在其进入扩展模式时可以修正from_u的path_cost
        return self.open_min()

    def _modify_cost2(self, to_v):
        nbs = self.neighbors(to_v)
        for nb in nbs:
            self.open_enqueue(nb)
        for dy in self.dynamic:
            dy_state = self.graph[dy]
            #此处省略视距检查
            if dy_state.f == "CLOSED":
                self.open_enqueue(dy_state)
        return self.open_min()

    def _init(self):
        self.init_graph()
        self.graph[self.goal].h = 0
        self.graph[self.goal].k = 0
        self.open = PriorityQueue(identby='pos', sortby='k')
        print("goal", self.goal)
        goal_s = self.graph[3, 135]
        print(goal_s)
        # print(goal_s.pos, goal_s.name)
        self.open_enqueue(self.graph[self.goal])

    @property
    def paths(self):
        if hasattr(self, "_paths"):
            return self._paths
        else:
            return None

    @paths.setter
    def paths(self, paths):
        self._paths = paths

    def add_path(self, path):
        if not self.paths:
            self.paths = []
        self.paths.append(path)

    def _generate_path(self, from_pos, to_pos=None):
        cache = []
        if to_pos is None:
            to_pos = self.goal
        cur_state = self.graph[from_pos]
        print(type(cur_state), cur_state)
        while cur_state and cur_state.pos != to_pos:
            cache.append(cur_state.pos)
            cur_state = cur_state.bp
        if cur_state:
            cache.append(cur_state.pos)
        self.add_path(cache)

    def dynamic_detector(self, center, radius=20):
        detected = []
        state = center
        while state:
            dist = self.block_dist(center, state)
            if dist <= radius/2 and self.has_dynamic(state):
                detected.append(state)
            state = state.bp
        if detected:
            for dy in self.dynamic:
                dy_state = self.graph[dy]
                dist = self.block_dist(center, dy_state)
                if dist <= radius and dy_state not in detected:
                    detected.append(dy_state)
        return detected

    def search(self, start=None, goal=None): #search
        if start is not None:
            self.start = start 
        if goal is not None:
            self.goal = goal
        #初始化
        self._init() 
        #首次规划路径
        while self.open_min() and self.open_min().pos != self.start: #open_min None or state
            self._process_state()
        if self.open_min().pos != self.start:#未到达起点
            print("unreached start point, break!")
            return False
        print("reach start: {}".format(self.open.head))
        self._generate_path(self.start) #存储路径
        self.enable_dynamic() #设定障碍物 非源程序逻辑，此处作为障碍物突然出现的一个trick手法。
        # 把起点移出队列
        state = self.open_dequeue()
        print("move from: {}".format(state.name))
        times = 0
        while True: #无限循环 python 不支持 do……while语法，用while True …… break来实现。
            if times == 10:
                raise Exception()
            times += 1
            # next_state = next_state.bp
            print("cur_state before move: {}".format(state))
            # while state.pos != self.goal and not self.has_dynamic(state.bp):
            while state.pos != self.goal and not self.dynamic_detector(state):
                #move to state.bp
                print("move to: {}".format(state.bp.name))
                state = state.bp
            if state.pos == self.goal: #到达终点
                print("reach goal: {}".format(state.name))
                return True
            print("cur_state after move: {}".format(state))
            if self.has_dynamic(state.bp):
                print("find dynamic: {}".format(state.bp))
            #state.bp有dynamic, 实际将state.bp加入队列，处于LOWER模式，扩展时将state加入队列。
            # print("queue head before modify: {}".format(self.open_min()))
            # self._modify_cost(state.bp, state, self.cost_dynamic) 
            # self._batch_modify_cost(state.bp, state, self.cost_dynamic)  #
            # self._batch_modify_visible_cost(state)
            dyns = self.dynamic_detector(state)
            print("dyns detected: {}".format(dyns))
            # raise Exception()
            if dyns:
                self._batch_modify_detected_cost(dyns)
            #把视距范围内的障碍物都要修改掉
            # self._modify_cost2(state.bp) 
            #另外的做法，将state.bp和8个neighbor加入队列，其中state.bp.bp的扩展用来修正state.pb的h值
            print("queue head after modify: {}".format(self.open_min()))
            while self.open_min() and self.open_min().k < state.h:
                self._process_state()
            print("queue head after process_state: {}".format(self.open_min()))
            print("cur state after process_state: {}".format(state))
            print("next state after process_state: {}".format(state.bp))
            nbs = self.neighbors(state.bp)
            for nb in nbs:
                print("-- bp.nb: {}".format(nb))
            # raise Exception()
            if not self.open_min():
                print("can't find new path, break")
                # raise Exception()
                return False
            # raise Exception()
            self._generate_path(state.pos) #存储路径

    def show_basemap(self):
        plt.axis('off')
        plt.gcf().set_size_inches(12, 8)
        plt.imshow(self.basemap, cmap='Greys_r')

    def draw_points(self):
        plt.plot(self.start[1], self.start[0], 'bo', color='green', markersize=12)
        plt.plot(self.goal[1], self.goal[0], 'bo', color='red', markersize=12)
        if hasattr(self, "dynamic"):
            p_x = []
            p_y = []
            if self.dynamic:
                for pos in self.dynamic:
                    p_x.append(pos[0])
                    p_y.append(pos[1])
                plt.plot(p_y, p_x, 'bo', color='gray', markersize=4)

    def draw_paths(self):
        colors = ["blue", "cyan", "magenta", "yellow"]
        for i in range(len(self.paths)):
            print("draw path: {}".format(i))
            print("path:{}\t{}".format(i, self.paths[i]))
            p_x = []
            p_y = []
            for pos in self.paths[i]:
                p_x.append(pos[0])
                p_y.append(pos[1])
            plt.plot(p_y, p_x, 'bo', color=colors[int(i%len(colors))], markersize=4)

    def show(self):
        plt.gcf().set_figheight(6)
        plt.gcf().set_figwidth(8)
        self.show_basemap()
        self.draw_points()
        self.draw_paths()
        plt.show()
    
    def print_nbs_by_pos(self, pos):
        state = self.graph[pos]
        return self.print_nbs(state)

    def print_nbs(self, state):
        nbs = self.neighbors(state)
        print("center: {}".format(state))
        for nb in nbs:
            print("nb: {}".format(nb))

def test():
    map_file = "./road.png"
    start = (20,15)
    goal = (130, 135)
    dynamic = [(40,117), (40,118), (40,119), (40,120), (40,121), (40,122), (40, 123), (40, 124), (40, 125), (40, 126), (40, 127), (40, 128), (40, 129), (40, 130)]
    dynamic = [(40,117), (40,118), (40,119), (40,120), (40,121), (40,122), (40, 123), (40, 124), (40, 125), (40, 126), (40, 127), (40, 128), (40, 129)]
    dstar_finder = DStarFinder(start=start, goal=goal, map_file=map_file, dynamic=dynamic)
    dstar_finder.search()
    # dstar_finder.print_nbs_by_pos((38, 118))
    dstar_finder.show()
    

if __name__ == "__main__":
	test()